import os
import argparse
import torch
import numpy as np
import open3d as o3d
from PIL import Image
import sys

# ensure repository root is on sys.path so nearby packages (zmqserver, zmqclient) can be imported
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if REPO_ROOT not in sys.path:
    sys.path.insert(0, REPO_ROOT)

from gsnet import AnyGrasp
from graspnetAPI import GraspGroup
from zmqserver.server import ZMQServer
from yolo11s.yolo_wrapper import YoloSegmentor

parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint_path', required=True, help='Model checkpoint path')
parser.add_argument('--yolo_model', type=str, default='/home/spaceai/anygrasp_sdk/yolo11s/log/yolo11s-seg.pt', help='Path to YOLO segmentation model (optional)')
parser.add_argument('--target_class', type=str, default='bottle', help='Target class name for segmentation (e.g., "bottle")')
parser.add_argument('--max_gripper_width', type=float, default=0.08, help='Maximum gripper width (<=0.1m)')
parser.add_argument('--gripper_height', type=float, default=0.03, help='Gripper height')
parser.add_argument('--top_down_grasp', action='store_true', help='Output top-down grasps.')
parser.add_argument('--debug', action='store_true', help='Enable debug mode')
cfgs = parser.parse_args()
cfgs.max_gripper_width = max(0, min(0.1, cfgs.max_gripper_width))

def demo():
    anygrasp = AnyGrasp(cfgs)
    anygrasp.load_net()
    # get camera intrinsics (same as below)
    fx, fy = 621.62, 621.62
    cx, cy = 640.00, 363.00
    scale = 1000.0

    # set workspace to filter output grasps (same as below)
    xmin, xmax = -0.19, 0.12
    ymin, ymax = 0.02, 0.15
    zmin, zmax = 0.0, 1.0
    lims = [xmin, xmax, ymin, ymax, zmin, zmax]

    # Run as ZMQ REP server: receive images from robot (REQ), run inference and reply.
    # Optionally load YOLO segmentor if requested (imports are at module top).
    yolo_model_path = getattr(cfgs, 'yolo_model', '')
    yolo_seg = None
    if yolo_model_path:
        yolo_seg = YoloSegmentor(yolo_model_path)
        print(f'Loaded YOLO segmentor: {yolo_model_path}')

    server = ZMQServer(bind_addr='tcp://*:5555', anygrasp=anygrasp, cfgs=cfgs)
    print('ZMQ server enabled: entering request loop')
    # main loop: get one observation, run same processing as below, reply result
    while True:
        header, color_img, depth_img = server.get_obs()

        # If a YOLO segmentor is configured, mask depth to keep only target class
        if yolo_seg is not None:
            # when debug is enabled, save an overlay visualization for inspection
            save_overlay = None
            if cfgs.debug:
                save_overlay = os.path.join(REPO_ROOT, 'yolo11s', 'output', 'last_yolo_overlay.png')
            depth_img = yolo_seg.mask_depth(color_img, depth_img, target=cfgs.target_class, save_overlay_path=save_overlay)

        # color_img: HxWx3 float32 [0,1], depth_img: HxW (raw units)
        colors = color_img.astype(np.float32)
        depths = np.array(depth_img)

        # get point cloud
        xmap, ymap = np.arange(depths.shape[1]), np.arange(depths.shape[0])
        xmap, ymap = np.meshgrid(xmap, ymap)
        points_z = depths / scale
        points_x = (xmap - cx) / fx * points_z
        points_y = (ymap - cy) / fy * points_z

        mask = (points_z > 0) & (points_z < 1)
        points = np.stack([points_x, points_y, points_z], axis=-1)
        points = points[mask].astype(np.float32)
        colors_masked = colors[mask].astype(np.float32)
        # Guard against empty point clouds (can happen after YOLO masking)
        if points.size == 0:
            print('No valid 3D points after masking/depth filtering!')
            # reply with a clear reason so client doesn't hang waiting
            try:
                server.reply({'ok': False, 'reason': 'no valid 3D points after masking/depth filtering'})
            except Exception:
                pass
            continue
        print(points.min(axis=0), points.max(axis=0))

        # Run inference and handle errors/empty results gracefully so server doesn't exit
        out = anygrasp.get_grasp(points, colors_masked, lims=lims, apply_object_mask=True, dense_grasp=False, collision_detection=True)

        if not out:
            print('AnyGrasp returned no result (None or empty)')
            server.reply({'ok': False, 'reason': 'no grasp returned'})
            continue

        gg, cloud = out

        # gg may be None or an empty collection; guard against None and len() failures
        try:
            empty = (gg is None) or (len(gg) == 0)
        except Exception:
            empty = True

        if empty:
            print('No Grasp detected after collision detection!')
            server.reply({'ok': False, 'reason': 'no grasp after collision detection'})
            continue

        gg = gg.nms().sort_by_score()

        gg_pick = gg[0:20]
        print(gg_pick.scores)
        print('grasp score:', gg_pick[0].score)

        best = gg_pick[0]

        # Robust extraction using graspnetAPI.grasp.Grasp fields
        # Preferred fields (based on dir): 'translation', 'rotation_matrix', 'width'
        def _to_list_safe(x):
            try:
                return np.array(x).tolist()
            except Exception:
                return None

        translation = None
        rotation = None
        width = None

        # translation: directly available
        if hasattr(best, 'translation'):
            translation = _to_list_safe(best.translation)

        # rotation: grasp object exposes 'rotation_matrix'
        if hasattr(best, 'rotation_matrix'):
            rotation = _to_list_safe(best.rotation_matrix)
        elif hasattr(best, 'rotation'):
            # fallback if attribute name differs
            rotation = _to_list_safe(best.rotation)

        # gripper width: try common names; 'width' exists on this object
        for name in ('width', 'jaw_width', 'gripper_width'):
            if hasattr(best, name):
                try:
                    width = float(getattr(best, name))
                except Exception:
                    width = None
                break

        # debug print extracted values
        try:
            print(f"extracted translation={translation}")
            print(f"extracted rotation_matrix={rotation}")
            print(f"extracted width={width}")
        except Exception:
            pass
        score = float(getattr(best, 'score', 0.0))

        if translation is None or rotation is None:
            server.reply({'ok': False, 'reason': 'could not extract pose from grasp object'})
            continue

        reply = {
            'ok': True,
            'translation': translation,
            'rotation': rotation,
            'gripper_width': width,
            'score': score,
        }

        # core reply contains pose and score only

        # if debug mode requested, inform client to shutdown and then visualize locally once
        if cfgs.debug:
            reply['shutdown'] = True

        # send reply to client
        server.reply(reply)

        # visualization (if debug) — show and then break loop (single-shot)
        if cfgs.debug:
            trans_mat = np.array([[1,0,0,0],[0,-1,0,0],[0,0,-1,0],[0,0,0,1]])
            try:
                cloud.transform(trans_mat)
            except Exception:
                pass
            grippers = gg.to_open3d_geometry_list()
            for gripper in grippers:
                try:
                    gripper.transform(trans_mat)
                except Exception:
                    pass
            o3d.visualization.draw_geometries([*grippers, cloud])
            o3d.visualization.draw_geometries([grippers[0], cloud])
            print('Debug visualization closed: exiting server loop')
            break
        # continue waiting for the next request (no return here)

if __name__ == '__main__':
    
    demo()
